#!/usr/bin/env python
# edit by yzl 

import numpy as np
from pycrazyswarm import *


Z = 0.7
d = 0.3
sleepRate = 30


def goformation(timeHelper, cfs, totalTime, radius, kPosition):
        #for cf, v in enumerate(cfs):
        #    startPos = cf.initialPosition + np.array([0, 0, Z])
        # r: virtual leader
        r = [0, 0, Z] 

        timeStage1 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage1
            if time > totalTime/4:
                break
            #r[2] = time / totalTime * 0.5 + Z
            #vr = [0, 0, 0.5 / totalTime]
            vr = [0, 0, 0]
            for i,cf in enumerate(cfs):
                print('stage 1 uav {}'.format(i))
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega0), np.sin(omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        totalTime = totalTime / 2

        # Stage2
        # rise up  to 1m height
        timeStage2 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage2
            if time > 2 * totalTime:
                break
            r[2] = time / (2 * totalTime) * 0.5 + Z
            vr = [0, 0, 0.5 / (2 * totalTime)]
            #vr = [0, 0, 0]
            for i,cf in enumerate(cfs):
                print('stage 2 uav {}'.format(i))
                omega = 2 * np.pi / totalTime
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = -radius * omega * np.sin(omega * time + omega0)
                vy = radius * omega * np.cos(omega * time + omega0)
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)
        
        
        # Stage3
        # fall down to 0.5m height
        z0 = r[2]
        #timeHelper.sleepForRate(sleepRate)
        timeStage3 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage3
            if time > 2 * totalTime:
                break
            r[2] = -time / (2 * totalTime) * 0.5 + z0
            vr = [0, 0, -0.5 / (2 * totalTime)]
            for i,cf in enumerate(cfs):
                print('stage 3 uav {}'.format(i))
                omega = 2 * np.pi / totalTime
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = -radius * omega * np.sin(omega * time + omega0)
                vy = radius * omega * np.cos(omega * time + omega0)
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)
        

        totalTime = totalTime / 3
        
        
        # Stage4
        # square formation tracking center (0,0)->(d,d)
        #timeHelper.sleepForRate(sleepRate)
        timeStage4 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage4
            if time > totalTime:
                break
            r = [d * time / totalTime, d * time / totalTime, Z]
            vr = [d / totalTime, d / totalTime, 0]
            for i,cf in enumerate(cfs):
                print('stage 4 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy,0 ] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)
        
        # Stage5
        # square formation tracking center (d,d)->(d,-d)
        #timeHelper.sleepForRate(sleepRate)
        timeStage5 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage5
            if time > 2 * totalTime:
                break
            r[1] = d - d * time / totalTime
            vr = [0, -d / totalTime, 0]
            for i,cf in enumerate(cfs):
                print('stage 5 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy,0 ] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        # Stage6
        # square formation tracking center (d,-d)->(-d,-d)
        #timeHelper.sleepForRate(sleepRate)
        timeStage6 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage6
            if time > 2 * totalTime:
                break
            r[0] = d - d * time / totalTime
            vr = [-d / totalTime, 0, 0]
            for i,cf in enumerate(cfs):
                print('stage 6 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy,0 ] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        # Stage7
        # square formation tracking center (-d,-d)->(-d,d)
        #timeHelper.sleepForRate(sleepRate)
        timeStage7= timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage7
            if time > 2 * totalTime:
                break
            r[1] = -d + d * time / totalTime
            vr = [0, d / totalTime, 0]
            for i,cf in enumerate(cfs):
                print('stage 7 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy,0 ] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        # Stage8
        # square formation tracking center (-d,d)->(0,0)
        #timeHelper.sleepForRate(sleepRate)
        timeStage8 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage8
            if time > totalTime:
                break
            r[0] = -d + d * time / totalTime
            r[1] = d - d * time / totalTime
            vr = [d / totalTime, -d / totalTime, 0]
            for i,cf in enumerate(cfs):
                print('stage 8 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 4
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega * time + omega0), np.sin(omega * time + omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy,0 ] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)
        
        
        # Stage9
        # prepare for line tracking 
        #timeHelper.sleepForRate(sleepRate)
        timeStage9 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage9
            if time > totalTime:
                break
            r[0] = 0
            r[1] = 0
            vr = [0, 0, 0]
            for i,cf in enumerate(cfs):
                print('stage 9 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 7
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega0), np.sin(omega0), 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)



        # Stage10
        # track  a line
        #timeHelper.sleepForRate(sleepRate)
        timeStage10 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage10
            if time > totalTime:
                break
            r[0] = 0
            r[1] = 0
            vr = [0, 0, 0]
            for i,cf in enumerate(cfs):
                print('stage 10 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 8 
                vx = 0
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega0), 0, 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        totalTime = totalTime * 3
        # Stage11
        # expand the line external radius change from 0.5m to 1m
        #timeHelper.sleepForRate(sleepRate)
        timeStage11 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage11
            if time > totalTime:
                break
            r[0] = 0
            r[1] = 0
            vr = [0, 0, 0]
            radius = 0.5 + 0.5 * time / totalTime
            for i,cf in enumerate(cfs):
                print('stage 11 uav {}'.format(i))
                omega = 0
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 8
                vx = 0.5 / totalTime * np.cos(omega0)
                vy = 0
                desiredPos = np.array(r) + radius * np.array(
                    [np.cos(omega0), 0, 0])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)
        
        totalTime = totalTime / 3


        # Stage12
        center34 = (np.cos(np.pi / 7 - 4 * np.pi / 4) + np.cos(np.pi / 7 - 6 * np.pi / 4 )) / 2
        r34 = np.cos(np.pi / 7 - 6 * np.pi / 4) - center34
        # exchange the posion of 3 and 4
        #timeHelper.sleepForRate(sleepRate)
        timeStage12 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage12
            if time > totalTime:
                break
            r[0] = 0
            r[1] = 0
            vr = [0, 0, 0]
            radius = 1
            for i,cf in enumerate(cfs):
                print('stage 12 uav {}'.format(i))
                omega = np.pi / totalTime
                omega0 = -i * 2 * np.pi / len(cfs) + np.pi / 7
                if i == 2 :
                    desiredPos = np.array(r) + np.array([center34, 0, 0]) + r34 * np.array(
                        [np.cos(omega * time + np.pi), np.sin(omega * time + np.pi), 0])
                    vx = - r34 * omega * np.sin(omega * time + np.pi)
                    vy = r34 * omega * np.cos(omega * time + np.pi)
                    errorX = desiredPos - cf.position()
                elif i ==3 :
                    desiredPos = np.array(r) + np.array([center34, 0, 0]) + r34 * np.array(
                        [np.cos(omega * time), np.sin(omega * time), 0])
                    vx = - r34 * omega * np.sin(omega * time)
                    vy = r34 * omega * np.cos(omega * time)
                    errorX = desiredPos - cf.position()
                else:
                    vx = 0
                    vy = 0
                    desiredPos = np.array(r) + radius * np.array(
                        [np.cos(omega0), 0, 0])
                    errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, 0] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        
        

        # Stage13
        # sin in z direction
        #timeHelper.sleepForRate(sleepRate)
        timeStage13 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage13
            if time > 5 * totalTime:
                break
            r = [0, time / (5 * totalTime), Z]
            vr = [0, 1 / (5 * totalTime), 0]
            radius = 1
            for i,cf in enumerate(cfs):
                print('stage 13 uav {}'.format(i))
                omega = np.pi / (totalTime * 2)
                vx = 0
                vy = 0
                vz = 0.2 * omega * np.cos(omega * time)
                desiredPos = np.array(r) +  np.array(
                    [-0.4 * i + 0.8, 0, 0.2 * np.sin(omega * time)])
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, vz] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)


            
        # Stage14
        # form triangle formation
        #timeHelper.sleepForRate(sleepRate)
        timeStage14 = timeHelper.time()
        radius = 0.8
        h = np.array([[radius, 0 ,0.2],[0, 0, 0.2],[radius*np.cos(np.pi * 7 / 6), radius*np.sin(np.pi * 7 / 6), 0.2],
            [radius*np.cos(np.pi * 5 / 6), radius*np.sin(np.pi * 5 / 6), 0.2]])
        while True:
            time = timeHelper.time() - timeStage14
            if time > 2 * totalTime:
                break
            r = [0, 1, Z]
            vr = [0, 0, 0]
            for i,cf in enumerate(cfs):
                print('stage 14 uav {}'.format(i))
                omega = 0
                vx = 0
                vy = 0
                vz = 0
                desiredPos = np.array(r) +  h[i,...]
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array([vx, vy, vz] + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        # Stage15
        # triangle formation circle tracking
        #timeHelper.sleepForRate(sleepRate)
        timeStage15 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage15
            if time > 4 * totalTime:
                break
            omega = np.pi / (2 * totalTime)
            h = np.array([[radius*np.cos(omega * time), radius*np.sin(omega * time) ,0.2],[0, 0, 0.2],[radius*np.cos(np.pi * 7 / 6 + omega * time), radius*np.sin(np.pi * 7 / 6 + omega * time), 0.2],
                [radius*np.cos(np.pi * 5 / 6 + omega * time), radius*np.sin(np.pi * 5 / 6 + omega * time), 0.2]])
            hv = np.array([[-omega*radius*np.sin(omega * time), omega*radius*np.cos(omega * time) ,0],[0, 0, 0],[-omega*radius*np.sin(np.pi * 7 / 6 + omega * time), omega*radius*np.cos(np.pi * 7 / 6 + omega * time), 0],
                [-omega*radius*np.sin(np.pi * 5 / 6 + omega * time), omega*radius*np.cos(np.pi * 5 / 6 + omega * time), 0]])
            r = [radius * np.cos(omega * time + np.pi /2), radius * np.sin(omega * time + np.pi /2), Z]
            vr = [-omega * radius * np.sin(omega * time + np.pi /2), omega * radius * np.cos(omega * time + np.pi /2), 0]
            for i,cf in enumerate(cfs):
                print('stage 15 uav {}'.format(i))
                v = hv[i,...]
                desiredPos = np.array(r) +  h[i,...]
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array(v + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)
        # Stage16
        # triangle formation circle tracking circle radius from 1m to 0.5m 
        #timeHelper.sleepForRate(sleepRate)
        timeStage16 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage16
            if time > 4 * totalTime:
                break
            radius  = 0.8 - 0.4 * time / (4 * totalTime)
            omega = np.pi / (2 * totalTime)
            h = np.array([[radius*np.cos(omega * time), radius*np.sin(omega * time) ,0.2],[0, 0, 0.2],[radius*np.cos(np.pi * 7 / 6 + omega * time), radius*np.sin(np.pi * 7 / 6 + omega * time), 0.2],
                [radius*np.cos(np.pi * 5 / 6 + omega * time), radius*np.sin(np.pi * 5 / 6 + omega * time), 0.2]])
            hv = np.array([[-omega*radius*np.sin(omega * time), omega*radius*np.cos(omega * time) ,0],[0, 0, 0],[-omega*radius*np.sin(np.pi * 7 / 6 + omega * time), omega*radius*np.cos(np.pi * 7 / 6 + omega * time), 0],
                [-omega*radius*np.sin(np.pi * 5 / 6 + omega * time), omega*radius*np.cos(np.pi * 5 / 6 + omega * time), 0]])
            r = [radius * np.cos(omega * time + np.pi /2), radius * np.sin(omega * time + np.pi /2), Z]
            vr = [-omega * radius * np.sin(omega * time + np.pi /2) - 0.1 / totalTime * np.cos(omega * time + np.pi /2), 
                omega * radius * np.cos(omega * time + np.pi /2) - 0.1 / totalTime * radius * np.sin(omega * time + np.pi /2), 0]
            for i,cf in enumerate(cfs):
                print('stage 15 uav {}'.format(i))
                v = hv[i,...]
                desiredPos = np.array(r) +  h[i,...]
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array(v + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)


        # Stage17
        # back to origin 
        #timeHelper.sleepForRate(sleepRate)
        timeStage17 = timeHelper.time()
        while True:
            time = timeHelper.time() - timeStage17
            if time > 2 * totalTime:
                break
            omega = np.pi / (2 * totalTime)
            h = np.array([[radius*np.cos(omega * time), radius*np.sin(omega * time) ,0.2],[0, 0, 0.2],[radius*np.cos(np.pi * 7 / 6 + omega * time), radius*np.sin(np.pi * 7 / 6 + omega * time), 0.2],
                [radius*np.cos(np.pi * 5 / 6 + omega * time), radius*np.sin(np.pi * 5 / 6 + omega * time), 0.2]])
            hv = np.array([[-omega*radius*np.sin(omega * time), omega*radius*np.cos(omega * time) ,0],[0, 0, 0],[-omega*radius*np.sin(np.pi * 7 / 6 + omega * time), omega*radius*np.cos(np.pi * 7 / 6 + omega * time), 0],
                [-omega*radius*np.sin(np.pi * 5 / 6 + omega * time), omega*radius*np.cos(np.pi * 5 / 6 + omega * time), 0]])

            r = [0, 0.5 - 0.5 * time / (2 * totalTime), Z]
            vr = [0, -1/(4 * totalTime), 0]
            for i,cf in enumerate(cfs):
                print('stage 15 uav {}'.format(i))
                v = hv[i,...]
                desiredPos = np.array(r) +  h[i,...]
                errorX = desiredPos - cf.position()
                cf.cmdVelocityWorld(np.array(v + kPosition * errorX + np.array(vr)), yawRate=0)
                timeHelper.sleepForRate(sleepRate)

        # Stage18 wait for landing
        for cf in cfs:
            cf.notifySetpointsStop()
            cf.land(targetHeight=0.03, duration=2.0)
            # timeHelper.sleepForRate(sleepRate)
        timeHelper.sleep(2 + Z)
        
        
        
if __name__ == "__main__":
    swarm = Crazyswarm()
    #swarm = Crazyswarm(args="--sim --vis null")
    timeHelper = swarm.timeHelper
    allcfs = swarm.allcfs
    print('The number of uavs is ' + str(len(allcfs.crazyflies)))
    allcfs.takeoff(targetHeight=Z, duration=1.0+Z)
    timeHelper.sleep(2 + Z)
    goformation(timeHelper, allcfs.crazyflies, totalTime=30, radius=0.5, kPosition=0.5)
    # swarm.input.waitUntilButtonPressed()
    # allcfs.land(targetHeight=0.04, duration=2.0)
    
